import torch.nn as nn


class SimpleModel(nn.Module):
    def __init__(self, input_dim):
        super(SimpleModel, self).__init__()

        self.layers = nn.Sequential(
            nn.Linear(input_dim, 16),
            nn.Dropout(0.001),
            nn.LeakyReLU(),
            nn.Linear(16, 8),
            nn.LeakyReLU(),
            nn.Linear(8, 1)
        )
        self.criterion = nn.MSELoss(reduction='mean')

    def forward(self, x):
        x = self.layers(x)
        x = x.squeeze(1)
        return x

    def cal_loss(self, pred, label):
        return self.criterion(pred, label)